import pandas as pd
import ccxt
from datetime import datetime, timedelta
from typing import List, Dict, Optional

class HistoricalDataHandler:
    """
    历史数据处理器
    负责获取和管理历史市场数据
    """
    
    def __init__(
        self,
        exchange: str,
        symbols: List[str],
        start_date: str,
        end_date: str,
        timeframe: str = '1h'
    ):
        """
        初始化数据处理器
        
        Args:
            exchange: 交易所名称
            symbols: 交易对列表
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            timeframe: 时间周期
        """
        self.exchange = self._initialize_exchange(exchange)
        self.symbols = symbols
        self.start_date = datetime.strptime(start_date, '%Y-%m-%d')
        self.end_date = datetime.strptime(end_date, '%Y-%m-%d')
        self.timeframe = timeframe
        self.data = {}
        self.current_index = {}
        self.continue_backtest = True
        
        # 加载历史数据
        self._load_historical_data()
    
    def _initialize_exchange(self, exchange_name: str) -> ccxt.Exchange:
        """
        初始化交易所实例
        """
        try:
            exchange_class = getattr(ccxt, exchange_name)
            exchange = exchange_class({
                'enableRateLimit': True,
                'options': {'defaultType': 'spot'},
                'proxies': {
                    'http': 'http://127.0.0.1:7890',
                    'https': 'http://127.0.0.1:7890',
                }
            })
            return exchange
        except Exception as e:
            print(f"Exchange initialization error: {str(e)}")
            raise
    
    def _load_historical_data(self):
        """
        加载历史数据
        """
        try:
            for symbol in self.symbols:
                print(f"Loading historical data for {symbol}...")
                
                # 获取K线数据
                ohlcv = self.exchange.fetch_ohlcv(
                    symbol=symbol,
                    timeframe=self.timeframe,
                    since=int(self.start_date.timestamp() * 1000),
                    limit=1000
                )
                
                # 转换为DataFrame
                df = pd.DataFrame(
                    ohlcv,
                    columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                )
                
                # 转换时间戳
                df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                df.set_index('timestamp', inplace=True)
                
                # 过滤日期范围
                df = df[self.start_date:self.end_date]
                
                # 存储数据
                self.data[symbol] = df
                self.current_index[symbol] = 0
                
                print(f"Loaded {len(df)} candles for {symbol}")
                
        except Exception as e:
            print(f"Data loading error: {str(e)}")
            raise
    
    def get_latest_bars(self, N: int = 1) -> Optional[pd.DataFrame]:
        """
        获取最新的N个数据点
        
        Args:
            N: 获取的数据点数量
            
        Returns:
            包含N个最新数据点的DataFrame
        """
        try:
            latest_data = {}
            for symbol in self.symbols:
                if self.current_index[symbol] >= len(self.data[symbol]):
                    return None
                
                start_idx = max(0, self.current_index[symbol] - N + 1)
                end_idx = self.current_index[symbol] + 1
                latest_data[symbol] = self.data[symbol].iloc[start_idx:end_idx]
            
            # 如果只有一个交易对，直接返回其数据
            if len(self.symbols) == 1:
                return latest_data[self.symbols[0]]
            
            # 如果有多个交易对，合并数据
            return pd.concat(latest_data.values(), axis=1, keys=self.symbols)
            
        except Exception as e:
            print(f"Error getting latest bars: {str(e)}")
            return None
    
    def update_bars(self):
        """
        更新到下一个数据点
        """
        try:
            # 更新所有交易对的索引
            for symbol in self.symbols:
                self.current_index[symbol] += 1
                
                # 检查是否还有更多数据
                if self.current_index[symbol] >= len(self.data[symbol]):
                    self.continue_backtest = False
                    break
                    
        except Exception as e:
            print(f"Error updating bars: {str(e)}")
            self.continue_backtest = False
